package com.yootoo.extpmml.Entites;


import com.alibaba.fastjson.JSON;
import com.yootoo.extpmml.Entites.model.ModelDesc;
import com.yootoo.extpmml.Entites.ruleentites.Message;
import com.yootoo.extpmml.Entites.ruleentites.fact.CheckResult;
import com.yootoo.extpmml.smg3.DynamicStrategiesRunner;
import ml.dmlc.xgboost4j.java.Booster;
import ml.dmlc.xgboost4j.java.DMatrix;
import org.dmg.pmml.FieldName;
import org.jpmml.evaluator.Evaluator;
import org.jpmml.evaluator.FieldValue;
import org.jpmml.evaluator.InputField;
import org.jpmml.evaluator.TargetField;
import org.kie.api.KieBase;
import org.kie.api.runtime.KieSession;
import java.util.HashMap;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;

public class ModelEvaluator extends BaseEvaluator<ModelDesc>{
    private ModelDesc model;
    private Object innerEvaluator;

    public ModelEvaluator(ModelDesc model, Object innerEvaluator) {
        this.model = model;
        this.innerEvaluator = innerEvaluator;
    }

    public ModelEvaluator() {
    }

    public ModelDesc getModel() {
        return model;
    }

    public void setModel(ModelDesc model) {
        this.model = model;
    }

    public Object getInnerEvaluator() {
        return innerEvaluator;
    }

    public void setInnerEvaluator(Object innerEvaluator) {
        this.innerEvaluator = innerEvaluator;
    }

    public Map<String,Object> execute(Map<String,Object> inputVars, String[] featureNames) throws Exception {
        Map<String,Object> res = new HashMap<>();
        if (innerEvaluator instanceof Evaluator){
            Evaluator eva = (Evaluator) innerEvaluator;
            List<InputField> inputFields = eva.getInputFields();
            Map<FieldName, FieldValue> arguments = new LinkedHashMap<>();
            for (InputField inputField : inputFields){
                FieldName inputFieldName = inputField.getName();
                Object value = inputVars.get(inputFieldName.getValue());
                FieldValue inputFieldValue = inputField.prepare(value);
                arguments.put(inputFieldName,inputFieldValue);
            }
            Map<FieldName, ?> result = eva.evaluate(arguments);
            List<TargetField> targetFields = eva.getTargetFields();
            for (TargetField targetField : targetFields){
                FieldName targetFieldName = targetField.getName();
                Object targetFieldValue = result.get(targetFieldName);
                res.put(targetFieldName.getValue(),targetFieldValue);
            }
        }else if (innerEvaluator instanceof Booster){
            Booster booster = (Booster) innerEvaluator;
            float[] values = new float[featureNames.length];
            for (int i = 0; i < featureNames.length; ++i) {
                values[i] = Float.parseFloat(inputVars.get(featureNames[i]).toString());
            }
            DMatrix matrix = new DMatrix(values, 1, values.length, -9999999f);
            float[] predictedRes = booster.predict(matrix)[0];
            for (int j = 0;j < predictedRes.length; j++){
                res.put("label_"+j,predictedRes[j]);
            }
        }else if (innerEvaluator instanceof KieBase){
            KieBase kieBase = (KieBase) innerEvaluator;
            Message message = new Message();
            message.setInput(inputVars);
            CheckResult checkResult = new CheckResult();
            KieSession kieSession = kieBase.newKieSession();
            kieSession.insert(message);
            kieSession.insert(checkResult);
            kieSession.fireAllRules();
            res = JSON.parseObject(JSON.toJSONString(checkResult),HashMap.class);
            kieSession.dispose();
        }else if (innerEvaluator instanceof DynamicStrategiesRunner){
            DynamicStrategiesRunner strategiesRunner = (DynamicStrategiesRunner) innerEvaluator;
            res = strategiesRunner.runStrategies(inputVars);
        }else {
            throw new UnsupportedOperationException("暂时只支持Evaluator、Booster、Drools和Smg3");
        }
        return res;
    }
}
